{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torchvision\n",
    "import torchvision.transforms as transforms\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "from torch.autograd import Variable\n",
    "import torch.optim as optim\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "from PIL import Image\n",
    "\n",
    "\n",
    "def loadtraindata():\n",
    "    path = r\"./train\"  # 路径\n",
    "    trainset = torchvision.datasets.ImageFolder(path,\n",
    "                                                transform=transforms.Compose([\n",
    "                                                    transforms.Resize((100, 100)),  # 将图片缩放到指定大小（h,w）\n",
    "\n",
    "                                                    transforms.CenterCrop(32),\n",
    "                                                    transforms.ToTensor()])\n",
    "                                                )\n",
    "\n",
    "    trainloader = torch.utils.data.DataLoader(trainset, batch_size=4,\n",
    "                                              shuffle=True, num_workers=2)\n",
    "    return trainloader\n",
    "\n",
    "\n",
    "class Net(nn.Module):  \n",
    "\n",
    "    def __init__(self):\n",
    "        super(Net, self).__init__()\n",
    "        self.conv1 = nn.Conv2d(3, 6, 5)  # 卷积层\n",
    "        self.pool = nn.MaxPool2d(2, 2)  # 池化层\n",
    "        self.conv2 = nn.Conv2d(6, 16, 5)  # 卷积层\n",
    "        self.fc1 = nn.Linear(400, 120)  # 全连接层\n",
    "        self.fc2 = nn.Linear(120, 84)\n",
    "        self.fc3 = nn.Linear(84, 2)  # 输出\n",
    "\n",
    "    def forward(self, x):  # 前向传播\n",
    "\n",
    "        x = self.pool(F.relu(self.conv1(x)))  \n",
    "        x = self.pool(F.relu(self.conv2(x)))\n",
    "        x = x.view(-1, 400)  # .view( )是一个tensor的方法，使得tensor改变size但是元素的总数是不变的。\n",
    "        # 从卷基层到全连接层的转换\n",
    "\n",
    "        x = F.relu(self.fc1(x))\n",
    "        x = F.relu(self.fc2(x))\n",
    "        x = self.fc3(x)\n",
    "        return x\n",
    "\n",
    "classes = ('不是奥特曼', '奥特曼')\n",
    "\n",
    "\n",
    "def loadtestdata():\n",
    "    path = r\"./test\"  # 测试集\n",
    "    testset = torchvision.datasets.ImageFolder(path,\n",
    "                                               transform=transforms.Compose([\n",
    "                                                   transforms.Resize((100, 100)),\n",
    "                                                   transforms.ToTensor()])\n",
    "                                               )\n",
    "    testloader = torch.utils.data.DataLoader(testset, batch_size=25,\n",
    "                                             shuffle=True, num_workers=2)\n",
    "    return testloader\n",
    "\n",
    "\n",
    "def trainandsave():  # 训练\n",
    "    trainloader = loadtraindata()\n",
    "    # 神经网络结构\n",
    "    net = Net()\n",
    "    optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)\n",
    "    criterion = nn.CrossEntropyLoss()  # 损失函数\n",
    "    # 训练部分\n",
    "    for epoch in range(1000):  # 训练的数据量为5个epoch，每个epoch为一个循环\n",
    "        running_loss = 0.0  # 定义一个变量方便对loss进行输出\n",
    "        for i, data in enumerate(trainloader, 0):  # 这里我们遇到了第一步中出现的trailoader，代码传入数据\n",
    "            inputs, labels = data  # data是从enumerate返回的data，包含数据和标签信息，分别赋值给inputs和labels\n",
    "            inputs, labels = Variable(inputs), Variable(labels)  # 转换数据格式用Variable\n",
    "\n",
    "            optimizer.zero_grad()  # 梯度置零\n",
    "            outputs = net(inputs)  # 把数据输进CNN网络net\n",
    "            loss = criterion(outputs, labels)  # 计算损失值\n",
    "            loss.backward()  # loss反向传播\n",
    "            optimizer.step()  # 反向传播后参数更新\n",
    "            running_loss += loss.item()  # loss累加\n",
    "            if i % 40 == 39:\n",
    "                print('[%d, %5d] loss: %.3f' % (epoch + 1, i + 1, running_loss / 40))  # 然后再除以200，就得到这两百次的平均损失值\n",
    "                running_loss = 0.0  # 这一个200次结束后，就把running_loss归零，下一个200次继续使用\n",
    "\n",
    "    print('Finished Training')\n",
    "    # 保存神经网络\n",
    "    torch.save(net, 'net.pkl')  # 保存结构和模型参数\n",
    "    torch.save(net.state_dict(), 'net_params.pkl')  # 只保存神经网络的模型参数\n",
    "\n",
    "\n",
    "def reload_net():\n",
    "    trainednet = torch.load('net.pkl')\n",
    "    return trainednet\n",
    "\n",
    "\n",
    "def imshow(img):\n",
    "    img = img / 2 + 0.5  \n",
    "    npimg = img.numpy()\n",
    "    plt.imshow(np.transpose(npimg, (1, 2, 0)))\n",
    "    plt.show()\n",
    "\n",
    "\n",
    "def test():  # 测试\n",
    "    testloader = loadtestdata()\n",
    "    net = reload_net()\n",
    "    dataiter = iter(testloader)\n",
    "    images, labels = dataiter.next()\n",
    "    imshow(torchvision.utils.make_grid(images, nrow=5))  # nrow是每行显示的图片数量\n",
    "    print('GroundTruth: '\n",
    "          , \" \".join('%5s' % classes[labels[j]] for j in range(25)))  # 打印前25个测试图\n",
    "    outputs = net(Variable(images))\n",
    "    _, predicted = torch.max(outputs.data, 1)\n",
    "\n",
    "    print('Predicted: ', \" \".join('%5s' % classes[predicted[j]] for j in range(25)))\n",
    "\n",
    "\n",
    "def test_one(img_path):\n",
    "    model = reload_net()\n",
    "    transform_valid = transforms.Compose([\n",
    "        transforms.Resize((32, 32), interpolation=2),\n",
    "        transforms.ToTensor()\n",
    "    ]\n",
    "    )\n",
    "    img = Image.open(img_path)\n",
    "    img_ = transform_valid(img).unsqueeze(0)\n",
    "    outputs = model(img_)\n",
    "    _, indices = torch.max(outputs, 1)\n",
    "    result = classes[indices]\n",
    "    print('predicted:', result)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<enumerate at 0x20927f66440>"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "trainloader = loadtraindata()\n",
    "enumerate(trainloader, 0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "    for epoch in range(1):  # 训练的数据量为5个epoch，每个epoch为一个循环\n",
    "        running_loss = 0.0  # 定义一个变量方便对loss进行输出\n",
    "        for i, data in enumerate(trainloader, 0):  # 这里我们遇到了第一步中出现的trailoader，代码传入数据\n",
    "            inputs, labels = data  # data是从enumerate返回的data，包含数据和标签信息，分别赋值给inputs和labels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[tensor([[[[0.6196, 0.6000, 0.5765,  ..., 0.6235, 0.4078, 0.1490],\n",
       "           [0.6157, 0.5961, 0.5725,  ..., 0.4784, 0.2039, 0.1137],\n",
       "           [0.6157, 0.5922, 0.5765,  ..., 0.2784, 0.1020, 0.0941],\n",
       "           ...,\n",
       "           [0.0902, 0.0588, 0.0902,  ..., 0.1333, 0.0980, 0.0784],\n",
       "           [0.0902, 0.0863, 0.0941,  ..., 0.1333, 0.1255, 0.0941],\n",
       "           [0.0863, 0.1020, 0.0980,  ..., 0.1216, 0.1137, 0.0980]],\n",
       " \n",
       "          [[0.4784, 0.4627, 0.4549,  ..., 0.5490, 0.3451, 0.1098],\n",
       "           [0.4784, 0.4627, 0.4431,  ..., 0.4314, 0.1686, 0.0902],\n",
       "           [0.4863, 0.4627, 0.4392,  ..., 0.2627, 0.0980, 0.0902],\n",
       "           ...,\n",
       "           [0.0902, 0.1137, 0.0902,  ..., 0.0902, 0.1098, 0.1098],\n",
       "           [0.0824, 0.0941, 0.0980,  ..., 0.0941, 0.0980, 0.1020],\n",
       "           [0.0824, 0.0784, 0.0941,  ..., 0.1059, 0.0941, 0.1137]],\n",
       " \n",
       "          [[0.3686, 0.3569, 0.3451,  ..., 0.3882, 0.2588, 0.1098],\n",
       "           [0.3647, 0.3490, 0.3412,  ..., 0.3294, 0.1137, 0.0863],\n",
       "           [0.3647, 0.3451, 0.3412,  ..., 0.2353, 0.0784, 0.0745],\n",
       "           ...,\n",
       "           [0.0902, 0.0863, 0.0980,  ..., 0.1529, 0.1373, 0.1020],\n",
       "           [0.0941, 0.0980, 0.1059,  ..., 0.1176, 0.1216, 0.1176],\n",
       "           [0.0902, 0.0941, 0.1059,  ..., 0.1137, 0.1137, 0.1294]]],\n",
       " \n",
       " \n",
       "         [[[0.4000, 0.3294, 0.5176,  ..., 0.4706, 0.4314, 0.4314],\n",
       "           [0.3882, 0.2902, 0.4824,  ..., 0.4157, 0.5529, 0.5451],\n",
       "           [0.4039, 0.3098, 0.4431,  ..., 0.3373, 0.3804, 0.3490],\n",
       "           ...,\n",
       "           [0.7255, 0.7255, 0.7333,  ..., 0.5647, 0.5922, 0.9686],\n",
       "           [0.7216, 0.7255, 0.7216,  ..., 0.6706, 0.5569, 0.8941],\n",
       "           [0.7294, 0.7216, 0.7176,  ..., 0.7922, 0.7255, 0.6706]],\n",
       " \n",
       "          [[0.3529, 0.2667, 0.4980,  ..., 0.4039, 0.3725, 0.3647],\n",
       "           [0.3529, 0.2588, 0.4667,  ..., 0.3451, 0.4863, 0.4745],\n",
       "           [0.3922, 0.2824, 0.4196,  ..., 0.2667, 0.3098, 0.2784],\n",
       "           ...,\n",
       "           [0.3725, 0.3765, 0.3804,  ..., 0.5216, 0.5608, 0.9608],\n",
       "           [0.3686, 0.3725, 0.3725,  ..., 0.6588, 0.5333, 0.8824],\n",
       "           [0.3765, 0.3765, 0.3725,  ..., 0.7529, 0.6706, 0.6588]],\n",
       " \n",
       "          [[0.3765, 0.3333, 0.5608,  ..., 0.3922, 0.3490, 0.3412],\n",
       "           [0.4000, 0.3176, 0.5176,  ..., 0.3373, 0.4471, 0.4353],\n",
       "           [0.4588, 0.3843, 0.5059,  ..., 0.2667, 0.2980, 0.2706],\n",
       "           ...,\n",
       "           [0.2353, 0.2392, 0.2431,  ..., 0.5216, 0.5529, 0.9176],\n",
       "           [0.2275, 0.2353, 0.2353,  ..., 0.6353, 0.5216, 0.8431],\n",
       "           [0.2314, 0.2314, 0.2275,  ..., 0.6863, 0.6039, 0.6314]]]]),\n",
       " tensor([0, 0])]"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[[0.6196, 0.6000, 0.5765,  ..., 0.6235, 0.4078, 0.1490],\n",
       "          [0.6157, 0.5961, 0.5725,  ..., 0.4784, 0.2039, 0.1137],\n",
       "          [0.6157, 0.5922, 0.5765,  ..., 0.2784, 0.1020, 0.0941],\n",
       "          ...,\n",
       "          [0.0902, 0.0588, 0.0902,  ..., 0.1333, 0.0980, 0.0784],\n",
       "          [0.0902, 0.0863, 0.0941,  ..., 0.1333, 0.1255, 0.0941],\n",
       "          [0.0863, 0.1020, 0.0980,  ..., 0.1216, 0.1137, 0.0980]],\n",
       "\n",
       "         [[0.4784, 0.4627, 0.4549,  ..., 0.5490, 0.3451, 0.1098],\n",
       "          [0.4784, 0.4627, 0.4431,  ..., 0.4314, 0.1686, 0.0902],\n",
       "          [0.4863, 0.4627, 0.4392,  ..., 0.2627, 0.0980, 0.0902],\n",
       "          ...,\n",
       "          [0.0902, 0.1137, 0.0902,  ..., 0.0902, 0.1098, 0.1098],\n",
       "          [0.0824, 0.0941, 0.0980,  ..., 0.0941, 0.0980, 0.1020],\n",
       "          [0.0824, 0.0784, 0.0941,  ..., 0.1059, 0.0941, 0.1137]],\n",
       "\n",
       "         [[0.3686, 0.3569, 0.3451,  ..., 0.3882, 0.2588, 0.1098],\n",
       "          [0.3647, 0.3490, 0.3412,  ..., 0.3294, 0.1137, 0.0863],\n",
       "          [0.3647, 0.3451, 0.3412,  ..., 0.2353, 0.0784, 0.0745],\n",
       "          ...,\n",
       "          [0.0902, 0.0863, 0.0980,  ..., 0.1529, 0.1373, 0.1020],\n",
       "          [0.0941, 0.0980, 0.1059,  ..., 0.1176, 0.1216, 0.1176],\n",
       "          [0.0902, 0.0941, 0.1059,  ..., 0.1137, 0.1137, 0.1294]]],\n",
       "\n",
       "\n",
       "        [[[0.4000, 0.3294, 0.5176,  ..., 0.4706, 0.4314, 0.4314],\n",
       "          [0.3882, 0.2902, 0.4824,  ..., 0.4157, 0.5529, 0.5451],\n",
       "          [0.4039, 0.3098, 0.4431,  ..., 0.3373, 0.3804, 0.3490],\n",
       "          ...,\n",
       "          [0.7255, 0.7255, 0.7333,  ..., 0.5647, 0.5922, 0.9686],\n",
       "          [0.7216, 0.7255, 0.7216,  ..., 0.6706, 0.5569, 0.8941],\n",
       "          [0.7294, 0.7216, 0.7176,  ..., 0.7922, 0.7255, 0.6706]],\n",
       "\n",
       "         [[0.3529, 0.2667, 0.4980,  ..., 0.4039, 0.3725, 0.3647],\n",
       "          [0.3529, 0.2588, 0.4667,  ..., 0.3451, 0.4863, 0.4745],\n",
       "          [0.3922, 0.2824, 0.4196,  ..., 0.2667, 0.3098, 0.2784],\n",
       "          ...,\n",
       "          [0.3725, 0.3765, 0.3804,  ..., 0.5216, 0.5608, 0.9608],\n",
       "          [0.3686, 0.3725, 0.3725,  ..., 0.6588, 0.5333, 0.8824],\n",
       "          [0.3765, 0.3765, 0.3725,  ..., 0.7529, 0.6706, 0.6588]],\n",
       "\n",
       "         [[0.3765, 0.3333, 0.5608,  ..., 0.3922, 0.3490, 0.3412],\n",
       "          [0.4000, 0.3176, 0.5176,  ..., 0.3373, 0.4471, 0.4353],\n",
       "          [0.4588, 0.3843, 0.5059,  ..., 0.2667, 0.2980, 0.2706],\n",
       "          ...,\n",
       "          [0.2353, 0.2392, 0.2431,  ..., 0.5216, 0.5529, 0.9176],\n",
       "          [0.2275, 0.2353, 0.2353,  ..., 0.6353, 0.5216, 0.8431],\n",
       "          [0.2314, 0.2314, 0.2275,  ..., 0.6863, 0.6039, 0.6314]]]])"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "inputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([0, 0])"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "labels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<torch.utils.data.dataloader.DataLoader at 0x20927f5ac70>"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "loadtraindata()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<torch.utils.data.dataloader.DataLoader at 0x20927f5adf0>"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "trainloader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
